import numpy as np

from constants import RELEASE_MODE
from data_io import (M, N, T, demand_matrix, reachable_matrix,
                     site_bandwidth_arr, quantile_index)
from tools import logging_time_cost, QuantileUpdater, LargestKeeper

# 整体思路为初解 + 优化
# 初解生成方案：
#     *. 每个边缘节点有T * 0.05次不计费的大流量时刻，在该时刻边缘节点可以尽量的吸收流量
#     *. 先给每个边缘节点指定极端分配时刻（T * 0.05次不计费的大流量时刻）
#     *. 对所有M * T个客户端请求量全排序，按照从大到小的顺序分配流量，并顺便给边缘节点指定极端分配时刻
#     *. 当遍历到某个客户端某个时刻的流量请求，针对该客户端可达的服务器中，挑选一个能最大吸收流量的且极端分配次数有剩余的边缘节点
#     *. 将该边缘节点在该时刻标记为极端分配时刻，并尽量多的吸收流量，同时该边缘节点的极端分配次数减一
#     *. 如果该客户端连接的所有服务器的大流量次数都已经被分配，则该客户端将流量请求平均分配
# 优化方案：
#     *. 优化初解的0.95分位点的流量有两种思路：一种是直接降低0.95分位点时刻的流量；另外一种是把某个后0.05时刻的流量分配给同时刻其他节点使得流量降低到0.95分位点之前
#     *. 基于这两种思路，产生两个优化方案：一个是优化各个边缘节点的0.95分位点，另外一个是优化后0.05极端分配时刻
#     *. 优化各个边缘节点的0.95分位点:
#         *. 先获取初解的各个边缘节点的0.95分位点时刻，对于某个边缘节点，将0.95分位点时刻的流量分配到同时刻低0.95分位点和高0.05的边缘节点，
#         *. 直到该边缘节点的流量低于流量排序位于0.95分位点时刻的前一个位置的时刻的流量值，此时0.95分位点变为排序后的前一个位置的时刻的流量值，获得一个更优的解
#         *. 流量转移的过程中保持分数的单调降低（新解一定更好）
#     *. 优化后0.05极端分配时刻:
#         *. 前一种优化方案不会改变边缘节点的后0.05高流量时刻，但是初解选出的后0.05高流量时刻未必是最优的，有些高流量时刻可能吸收不了太多流量
#         *. 对于某个边缘节点，先从后0.05高流量时刻中挑选一个时刻，将该边缘节点在该时刻的流量分配给同时刻低0.95分位点和高0.05的边缘节点，
#         *. 如果流量转移后，该边缘节点的流量可以降到0.95分位点以下，则原0.95分位点自动变为后0.05极端分配时刻，调整了极端分配时刻，此时获得一个更优的解
#         *. 如果流量转移后，该边缘节点的流量无法降到0.95分位点以下，则维持该时刻原来的分配方案
#     *. 这两个优化交替进行，有奇效，即使初解效果不好，也能优化到一个非常好的分数

# 官方线下数据集初解分数：9637
# 官方线下数据集初解优化后分数：8689


class DemandLargestManager:
    """
    实时维护、更新客户端的最大流量请求
    """

    def __init__(self, demand_matrix_) -> None:
        self.demand_largest_keepers = [LargestKeeper(demand_matrix_[:, client_id]) for client_id in range(demand_matrix_.shape[1])]

    def get_largest_demand(self):
        """
        获取流量请求最大的（时刻、客户端id、流量大小）
        """
        largest_demand = -1.0
        largest_t = None
        largest_client_id = None
        for client_id, largest_updater in enumerate(self.demand_largest_keepers):
            largest_v, t = largest_updater.get()
            if largest_v > largest_demand:
                largest_demand = largest_v
                largest_t = t
                largest_client_id = client_id

        if largest_t is not None:
            self.demand_largest_keepers[largest_client_id].next()

        return largest_t, largest_client_id, largest_demand

    def put_back(self, t, client_id, demand):
        self.demand_largest_keepers[client_id].add(demand, t)


class DemandOrderAssigner:
    def __init__(self) -> None:
        self.release_mode = RELEASE_MODE  # 发布模式跳过一些调试步骤

        self.demand_matrix = demand_matrix
        self.site_bandwidth_arr = site_bandwidth_arr
        self.reachable_matrix = reachable_matrix

        self.M = M
        self.N = N
        self.T = T
        self.t = 0
        self.quantile_index = quantile_index

        self.demand_largest_manager = DemandLargestManager(self.demand_matrix)
        self.demand_lacked_matrix = demand_matrix.copy()
        # 站点极端分配次数计数
        self.site_big_bw_remained_count = np.ones(shape=self.N, dtype=np.int64) * (self.T - (self.quantile_index + 1))
        self.site_big_bw_remained_count *= (self.reachable_matrix.sum(axis=0) != 0)
        self.site_big_bw_remained_count = np.clip(self.site_big_bw_remained_count, a_min=0, a_max=None)
        # 标记边缘节点在该时刻是否是极端分配时刻
        self.site_big_bw_flag_matrix = np.zeros(shape=(self.T, self.N), dtype=np.int64)
        self.all_assign_matrix = None  # 三维矩阵：T * M * N

        self.uniform_assign_weight_matrix = self.reachable_matrix / self.reachable_matrix.sum(axis=1, keepdims=True)
        # 客户端连接强度
        self.client_connection_strength = self.reachable_matrix.sum(axis=1)

    def next_assign(self):
        assign_matrix = self.get_all_assign_matrix()[self.t]
        if not self.release_mode:
            self._check_assign_matrix(assign_matrix)
        self.t += 1
        return assign_matrix

    def current_score(self):
        if self.release_mode:
            return None
        site_bw_devoted_matrix = self.all_assign_matrix.sum(axis=1)
        site_bw_devoted_sorted_matrix = np.sort(site_bw_devoted_matrix, axis=0)
        return site_bw_devoted_sorted_matrix[self.quantile_index].sum()

    def get_all_assign_matrix(self):
        if self.all_assign_matrix is not None:
            return self.all_assign_matrix

        self.all_assign_matrix = np.zeros(shape=(self.T, self.M, self.N), dtype=np.int64)

        while self.site_big_bw_remained_count.any():
            t, client_id, largest_demand = self.demand_largest_manager.get_largest_demand()
            if t is None:
                break
            else:
                self._radical_assign(t, client_id, largest_demand)

        for t in range(self.T):
            self._fine_tune_radical_assign(t)
            self._uniform_assign(t)
            self._process_overflow(t)

        return self.all_assign_matrix

    def _check_assign_matrix(self, assign_matrix):
        assert (assign_matrix.sum(axis=1) == self.demand_matrix[self.t]).all()  # 满足client需要
        assert (assign_matrix.sum(axis=0) <= self.site_bandwidth_arr).all()  # site最大带宽限制
        assert (assign_matrix * (self.reachable_matrix == 0)).sum() == 0  # 不可达时，带宽分配为0
        assert (assign_matrix >= 0).all()  # 带宽大于等于0

    def _fine_tune_radical_assign(self, t):
        """
        把标记为大流量节点，但是流量还未超上限的打满
        """
        demand_lacked_arr = self.demand_lacked_matrix[t]
        if not demand_lacked_arr.any():
            return

        assign_matrix = self.all_assign_matrix[t]
        site_big_bw_flag_arr = self.site_big_bw_flag_matrix[t]
        if not site_big_bw_flag_arr.any():
            return

        site_bw_devoted_arr = assign_matrix.sum(axis=0)
        site_bw_remained_arr = self.site_bandwidth_arr - site_bw_devoted_arr

        site_remained_big_bw_arr = site_big_bw_flag_arr * site_bw_remained_arr
        if not site_remained_big_bw_arr.any():
            return

        for site_id in np.argwhere(site_remained_big_bw_arr).flatten():
            site_remained_big_bw = site_remained_big_bw_arr[site_id]
            reachable_client_arr = self.reachable_matrix[:, site_id]
            client_selected_weight_arr = reachable_client_arr * self.client_connection_strength
            for client_id in np.argsort(client_selected_weight_arr):  # 连接数越少，越优先选择
                client_selected_weight = client_selected_weight_arr[client_id]
                if not client_selected_weight:
                    continue
                client_demand_lacked = demand_lacked_arr[client_id]
                if not client_demand_lacked:
                    continue
                site_bw_devote = min(site_remained_big_bw, client_demand_lacked)
                assign_matrix[client_id, site_id] += site_bw_devote
                demand_lacked_arr[client_id] -= site_bw_devote
                site_remained_big_bw -= site_bw_devote
                if not site_remained_big_bw:
                    break

    def _uniform_assign(self, t):
        """
        均匀分配
        """
        assign_matrix = self.all_assign_matrix[t]
        demand_lacked_arr = self.demand_lacked_matrix[t]
        site_bw_devoted_arr = assign_matrix.sum(axis=0)
        site_bw_remained_arr = self.site_bandwidth_arr - site_bw_devoted_arr
        site_bw_remained_flag_arr = (site_bw_remained_arr != 0)

        for client_id in range(self.M):
            client_bw_lacked = demand_lacked_arr[client_id]
            if not client_bw_lacked:
                continue
            reachable_site_arr = self.reachable_matrix[client_id, :]
            selectable_site_arr = reachable_site_arr * site_bw_remained_flag_arr
            transfer_wb_arr = (client_bw_lacked / selectable_site_arr.sum()) * selectable_site_arr
            transfer_wb_arr = transfer_wb_arr.astype(np.int64)
            remained_little = client_bw_lacked - transfer_wb_arr.sum()
            _site_id = np.random.choice(self.N, p=selectable_site_arr/selectable_site_arr.sum())
            transfer_wb_arr[_site_id] += remained_little

            assign_matrix[client_id, :] += transfer_wb_arr
            demand_lacked_arr[client_id] = 0

    def _radical_assign(self, t, client_id,  largest_demand):
        """
        将大流量优先分配给极端分配节点，如果该客户端可达的边缘节点的极端分配次数全部用完，则将流量平均分配给可达边缘节点
        """
        assign_matrix = self.all_assign_matrix[t]
        demand_lacked_arr = self.demand_lacked_matrix[t]
        site_big_bw_flag_arr = self.site_big_bw_flag_matrix[t]

        site_bw_devoted_arr = assign_matrix.sum(axis=0)
        site_bw_remained_arr = self.site_bandwidth_arr - site_bw_devoted_arr

        reachable_site_arr = self.reachable_matrix[client_id, :]
        selectable_big_bw_site_arr = np.argwhere(reachable_site_arr * site_big_bw_flag_arr * (site_bw_remained_arr != 0)).flatten()
        if len(selectable_big_bw_site_arr):
            for site_id in selectable_big_bw_site_arr:
                site_bw_remained = site_bw_remained_arr[site_id]
                site_bw_devote = min(site_bw_remained, largest_demand)
                assign_matrix[client_id, site_id] += site_bw_devote
                demand_lacked_arr[client_id] -= site_bw_devote
                largest_demand -= site_bw_devote
                if not largest_demand:
                    return
            self.demand_largest_manager.put_back(t, client_id, largest_demand)
        else:
            site_max_bw_arr = demand_lacked_arr @ self.reachable_matrix
            selectable_site_max_bw_arr = reachable_site_arr * site_max_bw_arr
            selectable_site_max_bw_arr *= (self.site_big_bw_remained_count > 0)
            selectable_site_max_bw_arr *= (site_big_bw_flag_arr == 0)
            selectable_site_max_bw_arr = np.clip(selectable_site_max_bw_arr, a_min=None, a_max=site_bw_remained_arr)
            selected_site_id = np.argmax(selectable_site_max_bw_arr)
            if selectable_site_max_bw_arr[selected_site_id]:
                self.site_big_bw_remained_count[selected_site_id] -= 1
                self.site_big_bw_flag_matrix[t, selected_site_id] = 1

                site_bw_remained = site_bw_remained_arr[selected_site_id]
                site_bw_devote = min(site_bw_remained, largest_demand)

                assign_matrix[client_id, selected_site_id] += site_bw_devote
                demand_lacked_arr[client_id] -= site_bw_devote
                largest_demand -= site_bw_devote
                if largest_demand:
                    self.demand_largest_manager.put_back(t, client_id, largest_demand)

    def _process_overflow(self, t):
        """
        处理分配溢出问题，如果溢出，则将流量随机充分配给邻接边缘节点，持续这个过程，直到所有边缘节点都不溢出
        """
        assign_matrix = self.all_assign_matrix[t]
        site_bw_devoted_arr = assign_matrix.sum(axis=0)

        overflow_site_flag_arr = (site_bw_devoted_arr > self.site_bandwidth_arr)
        overflow_site_ids = np.argwhere(overflow_site_flag_arr).flatten()
        while len(overflow_site_ids):
            for o_site_id in overflow_site_ids:
                overflow_bw = site_bw_devoted_arr[o_site_id] - self.site_bandwidth_arr[o_site_id]
                reachable_client_arr = self.reachable_matrix[:, o_site_id]

                client_selected_weight = reachable_client_arr * self.client_connection_strength
                client_selected_weight *= (assign_matrix[:, o_site_id] != 0)
                client_id = np.random.choice(self.M, p=client_selected_weight/client_selected_weight.sum())

                transfer_client_reachable_site_arr = self.reachable_matrix[client_id, :]
                transfer_client_reachable_site_arr = transfer_client_reachable_site_arr * (~overflow_site_flag_arr)

                if not transfer_client_reachable_site_arr.any():
                    continue

                transfer_wb = min(assign_matrix[client_id, o_site_id], overflow_bw)
                transfer_wb_arr = (transfer_wb / transfer_client_reachable_site_arr.sum()) * transfer_client_reachable_site_arr
                transfer_wb_arr = transfer_wb_arr.astype(np.int64)
                remained_little = transfer_wb - transfer_wb_arr.sum()
                _site_id = np.random.choice(self.N, p=transfer_client_reachable_site_arr/transfer_client_reachable_site_arr.sum())
                transfer_wb_arr[_site_id] += remained_little

                assign_matrix[client_id, o_site_id] -= transfer_wb
                site_bw_devoted_arr[o_site_id] -= transfer_wb

                assign_matrix[client_id, :] += transfer_wb_arr
                site_bw_devoted_arr[:] += transfer_wb_arr

            overflow_site_flag_arr = (site_bw_devoted_arr > self.site_bandwidth_arr)
            overflow_site_ids = np.argwhere(overflow_site_flag_arr).flatten()

        return assign_matrix


class OptQuantile:
    """
    对初始解的各个边缘节点的0.95分位点的流量进行优化，将0.95分位点的流量分配到低0.95分位点和高0.05的边缘节点，
    流量转移的过程中保持分数的单调降低
    """

    def __init__(self,
                 initial_assigner,
                 optimize_unite=3000,
                 decay_rate=0.867,
                 slow_time_cost_rate=3,
                 time_out=180,
                 reassign_adjust_max_try=30,
                 reassign_upper_gap=10) -> None:
        self.release_mode = RELEASE_MODE  # 发布模式跳过一些调试步骤

        self.demand_matrix = demand_matrix
        self.site_bandwidth_arr = site_bandwidth_arr
        self.reachable_matrix = reachable_matrix

        self.M = M
        self.N = N
        self.T = T
        self.t = 0
        self.quantile_index = quantile_index

        self.initial_assigner = initial_assigner
        self.optimize_unite = optimize_unite
        self.decay_rate = decay_rate
        self.slow_time_cost_rate = slow_time_cost_rate
        self.time_out = time_out
        self.reassign_adjust_max_try = reassign_adjust_max_try
        self.reassign_upper_gap = reassign_upper_gap

        self.all_assign_matrix = None  # 三维矩阵：T * M * N
        self.site_bw_devoted_quantile_updater = None

    def next_assign(self):
        assign_matrix = self.get_all_assign_matrix()[self.t]
        if not self.release_mode:
            self._check_assign_matrix(assign_matrix)
        self.t += 1
        return assign_matrix

    def current_score(self):
        return self._site_quantile_arr().sum()

    def get_all_assign_matrix(self):
        if self.all_assign_matrix is not None:
            return self.all_assign_matrix

        self.all_assign_matrix = self.initial_assigner.get_all_assign_matrix()
        site_bw_devoted_matrix = self.all_assign_matrix.sum(axis=1)
        self.site_bw_devoted_quantile_updater = [QuantileUpdater(site_bw_devoted_matrix[:, site_id], self.quantile_index)
                                                 for site_id in range(self.N)]

        import time
        time_cost_list = []
        normal_time_cost = None
        start_time = time.time()

        opt_step = 0
        while True:
            opt_step += 1
            opt_success = False

            if not (opt_step % 500):
                print(f"opt_step: <{opt_step}> score after opt:{self.current_score()}")

            t1 = time.time()
            for site_id in np.random.permutation(self.N):
                opt_success = (opt_success or self._optimize_bw_quantile(site_id))
            t2 = time.time()
            time_cost = t2 - t1

            if opt_step < 1000:
                time_cost_list.append(time_cost)
            elif opt_step == 1000:
                normal_time_cost = np.max(time_cost_list)
            elif (time_cost / normal_time_cost) > self.slow_time_cost_rate:
                self.optimize_unite = max(int(self.optimize_unite * self.decay_rate), 1)
                print(f"decay optimize_unite to: {self.optimize_unite}")
                print(f"opt_step: <{opt_step}> score after opt:{self.current_score()}")

            if (t2 - start_time) > self.time_out:
                print("time out")
                break

            if not opt_success:
                break

        return self.all_assign_matrix

    def _optimize_bw_quantile(self, site_id):
        """
        优化该边缘节点的0.95分位点，将该站点在0.95分位点时刻的流量随机分配给其他处于0.95分位点以下或后0.05时刻的边缘节点，
        直到该边缘节点的流量低于排序后的0.95分位点的前一时刻
        """
        t, v = self.site_bw_devoted_quantile_updater[site_id].get()
        if not v:
            return False
        assign_matrix = self.all_assign_matrix[t].copy()
        site_bw_devoted_arr = assign_matrix.sum(axis=0)
        site_quantile_arr = self._site_quantile_arr()
        site_reassign_upper_arr = self._get_reassign_upper_arr(site_id, site_bw_devoted_arr, site_quantile_arr)

        overflow_site_ids = np.argwhere(site_bw_devoted_arr > site_reassign_upper_arr).flatten()
        retry_times = 0
        while len(overflow_site_ids) and (retry_times < self.reassign_adjust_max_try):
            for o_site_id in overflow_site_ids:
                site_reachable_arr = self.reachable_matrix[:, o_site_id]
                client_id = np.random.choice(self.M, p=site_reachable_arr/site_reachable_arr.sum())
                transfer_client_reachable_arr = self.reachable_matrix[client_id, :]
                transfer_site_selected_weight = transfer_client_reachable_arr * (site_reassign_upper_arr - site_bw_devoted_arr)
                transfer_site_selected_weight[o_site_id] = 0.0  # 原溢出站点不可选择
                transfer_site_id = np.argmax(transfer_site_selected_weight)
                # 条件成立说明该客户端只与当前溢出的站点可达，因此无法转移溢出流量
                if transfer_site_selected_weight[transfer_site_id] == 0.0:
                    # 跳出循环，等待下次重新找一个可以转移流量的客户端
                    continue
                overflow_bw = site_bw_devoted_arr[o_site_id] - site_reassign_upper_arr[o_site_id]
                transfer_wb = min(assign_matrix[client_id, o_site_id], overflow_bw)
                if not transfer_wb:
                    continue
                assign_matrix[client_id, o_site_id] -= transfer_wb
                site_bw_devoted_arr[o_site_id] -= transfer_wb
                assign_matrix[client_id, transfer_site_id] += transfer_wb
                site_bw_devoted_arr[transfer_site_id] += transfer_wb

            overflow_site_ids = np.argwhere(site_bw_devoted_arr > site_reassign_upper_arr).flatten()
            retry_times += 1

        opt_success = (not len(overflow_site_ids))
        if opt_success:
            self.all_assign_matrix[t] = assign_matrix
            for site_id, new_bw_devoted in enumerate(site_bw_devoted_arr):
                self.site_bw_devoted_quantile_updater[site_id].update(t, new_bw_devoted)
        return opt_success

    def _get_reassign_upper_arr(self, site_id, site_bw_devoted_arr, site_quantile_arr):
        """
        计算重分配时各个边缘节点的可分配上限
        """
        site_reassign_upper_arr = np.where(site_bw_devoted_arr > site_quantile_arr, self.site_bandwidth_arr, site_quantile_arr - self.reassign_upper_gap)
        pre_quantile_bw = self.site_bw_devoted_quantile_updater[site_id].get_v(-1)  # 0.95分位点前一位的时刻的分配量
        site_reassign_upper_arr[site_id] = (pre_quantile_bw - self.optimize_unite)
        site_reassign_upper_arr = np.clip(site_reassign_upper_arr, a_min=0, a_max=None)
        return site_reassign_upper_arr

    def _site_quantile_arr(self):
        quantile_arr = np.zeros(shape=self.N, dtype=np.int64)
        for site_id in range(self.N):
            quantile_arr[site_id] = self.site_bw_devoted_quantile_updater[site_id].get_v()
        return quantile_arr

    def _round_assign_matrix(self, assign_matrix):
        demand_arr = self.demand_matrix[self.t]
        assign_matrix = assign_matrix.astype(np.int64)
        site_bw_devoted_arr = assign_matrix.sum(axis=0)
        demand_lacked_arr = demand_arr - assign_matrix.sum(axis=1)

        for client_id in range(self.M):
            client_bw_lacked = demand_lacked_arr[client_id]
            if not client_bw_lacked:
                continue
            client_reachable_arr = self.reachable_matrix[client_id]
            selected_site_id = np.random.choice(self.N, p=client_reachable_arr/client_reachable_arr.sum())
            assign_matrix[client_id, selected_site_id] += client_bw_lacked
            site_bw_devoted_arr[selected_site_id] += client_bw_lacked

        overflow_site_ids = np.argwhere(site_bw_devoted_arr > self.site_bandwidth_arr).flatten()
        while len(overflow_site_ids):
            for o_site_id in overflow_site_ids:
                overflow_bw = site_bw_devoted_arr[o_site_id] - self.site_bandwidth_arr[o_site_id]
                site_reachable_arr = self.reachable_matrix[:, o_site_id]
                client_id = np.random.choice(self.M, p=site_reachable_arr/site_reachable_arr.sum())
                transfer_wb = min(assign_matrix[client_id, o_site_id], overflow_bw)
                transfer_client_reachable_arr = self.reachable_matrix[client_id, :]
                transfer_site_id = np.random.choice(self.N, p=transfer_client_reachable_arr/transfer_client_reachable_arr.sum())
                if transfer_site_id == o_site_id:
                    continue
                assign_matrix[client_id, o_site_id] -= transfer_wb
                site_bw_devoted_arr[o_site_id] -= transfer_wb
                assign_matrix[client_id, transfer_site_id] += transfer_wb
                site_bw_devoted_arr[transfer_site_id] += transfer_wb

            overflow_site_ids = np.argwhere(site_bw_devoted_arr > self.site_bandwidth_arr).flatten()

        return assign_matrix

    def _check_assign_matrix(self, assign_matrix):
        assert (assign_matrix.sum(axis=1) == self.demand_matrix[self.t]).all()  # 满足client需要
        assert (assign_matrix.sum(axis=0) <= self.site_bandwidth_arr).all()  # site最大带宽限制
        assert (assign_matrix * (self.reachable_matrix == 0)).sum() == 0  # 不可达时，带宽分配为0
        assert (assign_matrix >= 0).all()  # 带宽大于等于0


class OptRadicalTime:
    """
    优化极端分配时刻，尝试将边缘节点后0.05时刻的流量分配给其他节点
    如果流量转移后，该边缘节点的流量可以降到0.95分位点以下，则原0.95分位点自动变为后0.05极端分配时刻，此时获得一个更优的解，相当于调整了极端分配时刻
    如果流量转移后，该边缘节点的流量无法降到0.95分位点以下，则维持原来的分配方案
    """

    def __init__(self,
                 initial_assigner,
                 optimize_unite=4000,
                 decay_rate=0.867,
                 slow_time_cost_rate=3,
                 time_out=180,
                 reassign_adjust_max_try=30,
                 reassign_upper_gap=10) -> None:
        self.release_mode = RELEASE_MODE  # 发布模式跳过一些调试步骤

        self.demand_matrix = demand_matrix
        self.site_bandwidth_arr = site_bandwidth_arr
        self.reachable_matrix = reachable_matrix

        self.M = M
        self.N = N
        self.T = T
        self.t = 0
        self.quantile_index = quantile_index

        self.t_demand_sum_arr = self.demand_matrix.sum(axis=1)

        self.initial_assigner = initial_assigner
        self.optimize_unite = optimize_unite
        self.decay_rate = decay_rate
        self.slow_time_cost_rate = slow_time_cost_rate
        self.time_out = time_out
        self.reassign_adjust_max_try = reassign_adjust_max_try
        self.reassign_upper_gap = reassign_upper_gap

        self.all_assign_matrix = None  # 三维矩阵：T * M * N
        self.site_bw_devoted_quantile_updater = None

    def next_assign(self):
        assign_matrix = self.get_all_assign_matrix()[self.t]
        if not self.release_mode:
            self._check_assign_matrix(assign_matrix)
        self.t += 1
        return assign_matrix

    def current_score(self):
        return self._site_quantile_arr().sum()

    def get_all_assign_matrix(self):
        if self.all_assign_matrix is not None:
            return self.all_assign_matrix

        self.all_assign_matrix = self.initial_assigner.get_all_assign_matrix()
        site_bw_devoted_matrix = self.all_assign_matrix.sum(axis=1)
        self.site_bw_devoted_quantile_updater = [QuantileUpdater(site_bw_devoted_matrix[:, site_id], self.quantile_index)
                                                 for site_id in range(self.N)]

        import time
        time_cost_list = []
        normal_time_cost = None
        start_time = time.time()

        opt_step = 0
        while True:
            opt_step += 1
            opt_success = False

            if not (opt_step % 500):
                print(f"opt_step: <{opt_step}> score after opt:{self.current_score()}")

            t1 = time.time()
            for site_id in np.random.permutation(self.N):
                opt_success = (opt_success or self._optimize_bw_quantile(site_id))
            t2 = time.time()
            time_cost = t2 - t1

            if opt_step < 1000:
                time_cost_list.append(time_cost)
            elif opt_step == 1000:
                normal_time_cost = np.max(time_cost_list)
            elif (time_cost / normal_time_cost) > self.slow_time_cost_rate:
                self.optimize_unite = max(int(self.optimize_unite * self.decay_rate), 1)
                print(f"decay optimize_unite to: {self.optimize_unite}")
                print(f"opt_step: <{opt_step}> score after opt:{self.current_score()}")

            if (t2 - start_time) > self.time_out:
                print("time out")
                break

            if not opt_success:
                break

        return self.all_assign_matrix

    def _optimize_bw_quantile(self, site_id):
        t_arr, v_arr = self.site_bw_devoted_quantile_updater[site_id].get_all()  # 获取所有0.95分位点及之后的时间时刻
        if not v_arr[0]:
            return False

        # after_quantile_t_demand_sum_arr = self.t_demand_sum_arr[t_arr]
        # t = t_arr[np.argmin(after_quantile_t_demand_sum_arr)]
        t = np.random.choice(t_arr)
        t_demand_sum = self.t_demand_sum_arr[t]

        assign_matrix = self.all_assign_matrix[t].copy()
        site_bw_devoted_arr = assign_matrix.sum(axis=0)
        site_quantile_arr = self._site_quantile_arr()
        site_reassign_upper_arr = self._get_reassign_upper_arr(site_id, site_bw_devoted_arr, site_quantile_arr)

        if site_reassign_upper_arr.sum() < t_demand_sum:
            return False

        overflow_site_ids = np.argwhere(site_bw_devoted_arr > site_reassign_upper_arr).flatten()
        retry_times = 0
        while len(overflow_site_ids) and (retry_times < self.reassign_adjust_max_try):
            for o_site_id in overflow_site_ids:
                site_reachable_arr = self.reachable_matrix[:, o_site_id]
                client_id = np.random.choice(self.M, p=site_reachable_arr/site_reachable_arr.sum())
                transfer_client_reachable_arr = self.reachable_matrix[client_id, :]
                transfer_site_selected_weight = transfer_client_reachable_arr * (site_reassign_upper_arr - site_bw_devoted_arr)
                transfer_site_selected_weight[o_site_id] = 0.0  # 原溢出站点不可选择
                transfer_site_id = np.argmax(transfer_site_selected_weight)
                # 条件成立说明该客户端只与当前溢出的站点可达，因此无法转移溢出流量
                if transfer_site_selected_weight[transfer_site_id] == 0.0:
                    # 跳出循环，等待下次重新找一个可以转移流量的客户端
                    continue
                overflow_bw = site_bw_devoted_arr[o_site_id] - site_reassign_upper_arr[o_site_id]
                transfer_wb = min(assign_matrix[client_id, o_site_id], overflow_bw)
                if not transfer_wb:
                    continue
                assign_matrix[client_id, o_site_id] -= transfer_wb
                site_bw_devoted_arr[o_site_id] -= transfer_wb
                assign_matrix[client_id, transfer_site_id] += transfer_wb
                site_bw_devoted_arr[transfer_site_id] += transfer_wb

            overflow_site_ids = np.argwhere(site_bw_devoted_arr > site_reassign_upper_arr).flatten()
            retry_times += 1

        opt_success = (not len(overflow_site_ids))
        if opt_success:
            self.all_assign_matrix[t] = assign_matrix
            for site_id, new_bw_devoted in enumerate(site_bw_devoted_arr):
                self.site_bw_devoted_quantile_updater[site_id].update(t, new_bw_devoted)
        return opt_success

    def _get_reassign_upper_arr(self, site_id, site_bw_devoted_arr, site_quantile_arr):
        site_reassign_upper_arr = np.where(site_bw_devoted_arr > site_quantile_arr, self.site_bandwidth_arr, site_quantile_arr - self.reassign_upper_gap)
        pre_quantile_bw = self.site_bw_devoted_quantile_updater[site_id].get_v(-1)  # 0.95分位点前一位的时刻的分配量
        site_reassign_upper_arr[site_id] = (pre_quantile_bw - self.optimize_unite)
        site_reassign_upper_arr = np.clip(site_reassign_upper_arr, a_min=0, a_max=None)
        return site_reassign_upper_arr

    def _site_quantile_arr(self):
        quantile_arr = np.zeros(shape=self.N, dtype=np.int64)
        for site_id in range(self.N):
            quantile_arr[site_id] = self.site_bw_devoted_quantile_updater[site_id].get_v()
        return quantile_arr

    def _round_assign_matrix(self, assign_matrix):
        demand_arr = self.demand_matrix[self.t]
        assign_matrix = assign_matrix.astype(np.int64)
        site_bw_devoted_arr = assign_matrix.sum(axis=0)
        demand_lacked_arr = demand_arr - assign_matrix.sum(axis=1)

        for client_id in range(self.M):
            client_bw_lacked = demand_lacked_arr[client_id]
            if not client_bw_lacked:
                continue
            client_reachable_arr = self.reachable_matrix[client_id]
            selected_site_id = np.random.choice(self.N, p=client_reachable_arr/client_reachable_arr.sum())
            assign_matrix[client_id, selected_site_id] += client_bw_lacked
            site_bw_devoted_arr[selected_site_id] += client_bw_lacked

        overflow_site_ids = np.argwhere(site_bw_devoted_arr > self.site_bandwidth_arr).flatten()
        while len(overflow_site_ids):
            for o_site_id in overflow_site_ids:
                overflow_bw = site_bw_devoted_arr[o_site_id] - self.site_bandwidth_arr[o_site_id]
                site_reachable_arr = self.reachable_matrix[:, o_site_id]
                client_id = np.random.choice(self.M, p=site_reachable_arr/site_reachable_arr.sum())
                transfer_wb = min(assign_matrix[client_id, o_site_id], overflow_bw)
                transfer_client_reachable_arr = self.reachable_matrix[client_id, :]
                transfer_site_id = np.random.choice(self.N, p=transfer_client_reachable_arr/transfer_client_reachable_arr.sum())
                if transfer_site_id == o_site_id:
                    continue
                assign_matrix[client_id, o_site_id] -= transfer_wb
                site_bw_devoted_arr[o_site_id] -= transfer_wb
                assign_matrix[client_id, transfer_site_id] += transfer_wb
                site_bw_devoted_arr[transfer_site_id] += transfer_wb

            overflow_site_ids = np.argwhere(site_bw_devoted_arr > self.site_bandwidth_arr).flatten()

        return assign_matrix

    def _check_assign_matrix(self, assign_matrix):
        assert (assign_matrix.sum(axis=1) == self.demand_matrix[self.t]).all()  # 满足client需要
        assert (assign_matrix.sum(axis=0) <= self.site_bandwidth_arr).all()  # site最大带宽限制
        assert (assign_matrix * (self.reachable_matrix == 0)).sum() == 0  # 不可达时，带宽分配为0
        assert (assign_matrix >= 0).all()  # 带宽大于等于0


@logging_time_cost()
def main():
    initial_assigner_assigner = DemandOrderAssigner()
    initial_assigner_assigner.get_all_assign_matrix()
    print(f"initial_assigner score: {initial_assigner_assigner.current_score()}")

    opt_quantile = OptQuantile(initial_assigner_assigner,
                               optimize_unite=5000,
                               decay_rate=0.9,
                               slow_time_cost_rate=5,
                               time_out=30,
                               reassign_adjust_max_try=30,
                               reassign_upper_gap=10)
    opt_quantile.get_all_assign_matrix()
    print(f"opt_quantile score: {opt_quantile.current_score()}")

    max_stable_times = 5
    stable_times = 0
    last_score = 0
    while True:
        opt_radical_time = OptRadicalTime(opt_quantile,
                                          optimize_unite=2000,
                                          decay_rate=0.9,
                                          slow_time_cost_rate=5,
                                          time_out=12,
                                          reassign_adjust_max_try=30,
                                          reassign_upper_gap=10)
        opt_radical_time.get_all_assign_matrix()
        print(f"opt_radical_time score: {opt_radical_time.current_score()}")

        opt_quantile = OptQuantile(opt_radical_time,
                                   optimize_unite=2000,
                                   decay_rate=0.9,
                                   slow_time_cost_rate=5,
                                   time_out=30,
                                   reassign_adjust_max_try=30,
                                   reassign_upper_gap=10)
        opt_quantile.get_all_assign_matrix()
        print(f"opt_quantile score: {opt_quantile.current_score()}")

        new_score = opt_quantile.current_score()
        if last_score != new_score:
            last_score = new_score
            stable_times = 0
        elif stable_times < max_stable_times:
            stable_times += 1
        else:
            break

    if RELEASE_MODE:
        from data_io import output
        output(opt_quantile)
    else:
        from tools import eval
        eval(opt_quantile)


if __name__ == "__main__":
    main()
